import cnnbest_model
from load_ascad import *
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm

model_name = "cnn"
batch_size = 200
data_file_path = "ASCAD.h5"
epochs = 50
learn_rate = 0.00001
device = torch.device("cuda:0" if (torch.cuda.is_available()) else "cpu")
print(f"将使用{str(device)}作为设备")
# 加载数据集
train_data_loader, train_plaintext = load_ascad(data_file_path, batch_size, True, 0)
test_data_loader, test_plaintext = load_ascad(data_file_path, 1, False, 1)

if __name__ == "__main__":
    # 损失函数
    loss_function = nn.CrossEntropyLoss().to(device)
    loss_function.to(device)
    # 模型
    model = None
    if (model_name == "cnn"):
        model = cnnbest_model.CnnBest()
    else:
        print("错误：没有定义的model_name!")
        exit()
    model.to(device)
    # 优化器
    optimizer = optim.RMSprop(model.parameters(), lr=learn_rate)
    # 开始训练
    for epoch in range(epochs):
        print(f"epochs: {epoch}/{epochs}:")
        running_loss = 0.0
        train_right_count = 0.0
        for i, data in enumerate(tqdm(train_data_loader)):
            trace, label = data
            trace = trace.to(device)
            label = label.to(device)
            optimizer.zero_grad()
            outputs = model(trace)

            for index in range(batch_size):
                pred_label = torch.argmax(outputs, dim=1)[index].item()
                if (label[index].item() == pred_label):
                    train_right_count += 1
            loss = loss_function(outputs, label)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        # 训练过程的准确度
        train_acc = train_right_count / (len(train_data_loader) * batch_size)
        print(f" 平均损失: {running_loss/len(train_data_loader)} 训练数据集准确度：{train_acc}")
        if ((epoch + 1) == epochs):
            torch.save(model.state_dict(), f"./trained_model/{model_name}-last.pth")
